#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   evidence_GCN.py    
@Contact :   raogx.vip@hotmail.com
@License :   (C)Copyright 2020

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2021-04-08 9:58   WandongShi      1.0         None
'''
import torch
from torch import nn as nn
import torch.nn.functional as F

class evidence_GCN(nn.Module):
    def __init__(self, gcn_dim, num_layers):
        super(evidence_GCN, self).__init__()
        self.layers = num_layers
        self.in_dim = 768
        self.gcn_dim = gcn_dim
        self.gcn_drop = nn.Dropout(0.8)
        # gcn layer
        self.W = nn.ModuleList()
        for layer in range(self.layers):
            input_dim = self.in_dim if layer == 0 else self.gcn_dim
            self.W.append(nn.Linear(input_dim, self.gcn_dim))

    def forward(self, adj, inputs):
        adj = torch.from_numpy(adj)
        adj = adj.cuda(1)
        # gcn layer
        denom = adj.sum(2).unsqueeze(2) + 1  # norm
        for l in range(self.layers):
            Ax = adj.bmm(inputs)
            AxW = self.W[l](Ax)
            AxW = AxW / denom
            gAxW = F.relu(AxW)
            gcn_outputs = self.gcn_drop(gAxW) if l < self.layers - 1 else gAxW
        return gcn_outputs